/*
 * Copyright (c) HiSilicon (Shanghai) Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: 音频预处理模块 - 处理单声道数据，去除空声道，准备编码格式
 */

#include "little_test.h"
#include "osal_debug.h"
#include "securec.h"
#include <stdlib.h>
#include <string.h>

/**
 * @brief 初始化音频预处理模块
 * @param processing 音频预处理结构体指针
 * @return 0成功，其他失败
 */
int audio_processing_init(audio_processing_t *processing)
{
    if (processing == NULL) {
        osal_printk("[AUDIO_PROC] Invalid processing pointer\n");
        return -1;
    }

    // 初始化结构体
    memset_s(processing, sizeof(audio_processing_t), 0, sizeof(audio_processing_t));

    // 分配原始数据缓冲区（双声道，但只有一个声道有数据）
    processing->raw_buffer = (int16_t *)malloc(LITTLE_TEST_FRAME_SIZE * 2 * sizeof(int16_t));
    if (processing->raw_buffer == NULL) {
        osal_printk("[AUDIO_PROC] Failed to allocate raw buffer\n");
        return -1;
    }

    // 分配处理后的数据缓冲区（单声道）
    processing->processed_buffer = (int16_t *)malloc(LITTLE_TEST_FRAME_SIZE * sizeof(int16_t));
    if (processing->processed_buffer == NULL) {
        osal_printk("[AUDIO_PROC] Failed to allocate processed buffer\n");
        free(processing->raw_buffer);
        return -1;
    }

    // 分配临时缓冲区用于滤波等处理
    processing->temp_buffer = (int16_t *)malloc(LITTLE_TEST_FRAME_SIZE * sizeof(int16_t));
    if (processing->temp_buffer == NULL) {
        osal_printk("[AUDIO_PROC] Failed to allocate temp buffer\n");
        free(processing->raw_buffer);
        free(processing->processed_buffer);
        return -1;
    }

    // 初始化高通滤波器状态（用于去除直流分量）
    processing->hpf_prev_input = 0;
    processing->hpf_prev_output = 0;
    processing->hpf_alpha = 0.995f;  // 高通滤波器系数，截止频率约为80Hz

    // 初始化AGC（自动增益控制）参数
    processing->agc_target_level = 16384;  // 目标电平（约50%满量程）
    processing->agc_current_gain = 1.0f;
    processing->agc_attack_time = 0.01f;   // 攻击时间10ms
    processing->agc_release_time = 0.1f;   // 释放时间100ms
    processing->agc_max_gain = 8.0f;       // 最大增益
    processing->agc_min_gain = 0.125f;     // 最小增益

    // 初始化噪声门限
    processing->noise_gate_threshold = 500;  // 噪声门限
    processing->noise_gate_enabled = true;

    // 初始化统计信息
    processing->frame_count = 0;
    processing->peak_level = 0;
    processing->rms_level = 0;

    processing->initialized = true;
    
    osal_printk("[AUDIO_PROC] Initialized successfully\n");
    osal_printk("[AUDIO_PROC] Frame size: %d samples, HPF alpha: %.3f\n", 
                LITTLE_TEST_FRAME_SIZE, processing->hpf_alpha);
    osal_printk("[AUDIO_PROC] AGC target: %d, noise gate: %d\n", 
                processing->agc_target_level, processing->noise_gate_threshold);
    
    return 0;
}

/**
 * @brief 反初始化音频预处理模块
 * @param processing 音频预处理结构体指针
 */
void audio_processing_deinit(audio_processing_t *processing)
{
    if (processing == NULL) {
        return;
    }

    // 释放缓冲区
    if (processing->raw_buffer != NULL) {
        free(processing->raw_buffer);
        processing->raw_buffer = NULL;
    }

    if (processing->processed_buffer != NULL) {
        free(processing->processed_buffer);
        processing->processed_buffer = NULL;
    }

    if (processing->temp_buffer != NULL) {
        free(processing->temp_buffer);
        processing->temp_buffer = NULL;
    }

    processing->initialized = false;
    
    osal_printk("[AUDIO_PROC] Deinitialized\n");
}

/**
 * @brief 计算RMS电平
 * @param samples 音频样本数据
 * @param count 样本数量
 * @return RMS电平值
 */
static uint32_t calculate_rms_level(const int16_t *samples, uint32_t count)
{
    uint64_t sum_squares = 0;
    uint32_t i;
    
    for (i = 0; i < count; i++) {
        int32_t sample = samples[i];
        sum_squares += (uint64_t)(sample * sample);
    }
    
    // 计算均方根
    uint32_t mean_square = (uint32_t)(sum_squares / count);
    
    // 简单的平方根近似（使用位移）
    uint32_t rms = 0;
    uint32_t bit = 1 << 15;
    
    while (bit > mean_square) {
        bit >>= 2;
    }
    
    while (bit != 0) {
        if (mean_square >= rms + bit) {
            mean_square -= rms + bit;
            rms = (rms >> 1) + bit;
        } else {
            rms >>= 1;
        }
        bit >>= 2;
    }
    
    return rms;
}

/**
 * @brief 查找峰值电平
 * @param samples 音频样本数据
 * @param count 样本数量
 * @return 峰值电平值
 */
static uint32_t find_peak_level(const int16_t *samples, uint32_t count)
{
    uint32_t peak = 0;
    uint32_t i;
    
    for (i = 0; i < count; i++) {
        uint32_t abs_sample = (uint32_t)abs(samples[i]);
        if (abs_sample > peak) {
            peak = abs_sample;
        }
    }
    
    return peak;
}

/**
 * @brief 高通滤波器（去除直流分量）
 * @param processing 音频预处理结构体指针
 * @param input 输入样本
 * @return 滤波后的样本
 */
static int16_t high_pass_filter(audio_processing_t *processing, int16_t input)
{
    // 一阶高通滤波器：y[n] = alpha * (y[n-1] + x[n] - x[n-1])
    float output_f = processing->hpf_alpha * 
                    (processing->hpf_prev_output + input - processing->hpf_prev_input);
    
    // 限幅
    if (output_f > 32767.0f) {
        output_f = 32767.0f;
    } else if (output_f < -32768.0f) {
        output_f = -32768.0f;
    }
    
    int16_t output = (int16_t)output_f;
    
    // 更新状态
    processing->hpf_prev_input = input;
    processing->hpf_prev_output = output;
    
    return output;
}

/**
 * @brief 自动增益控制
 * @param processing 音频预处理结构体指针
 * @param samples 音频样本数据
 * @param count 样本数量
 */
static void apply_agc(audio_processing_t *processing, int16_t *samples, uint32_t count)
{
    // 计算当前帧的RMS电平
    uint32_t current_rms = calculate_rms_level(samples, count);
    
    if (current_rms > 0) {
        // 计算所需增益
        float target_gain = (float)processing->agc_target_level / (float)current_rms;
        
        // 限制增益范围
        if (target_gain > processing->agc_max_gain) {
            target_gain = processing->agc_max_gain;
        } else if (target_gain < processing->agc_min_gain) {
            target_gain = processing->agc_min_gain;
        }
        
        // 平滑增益变化
        float gain_diff = target_gain - processing->agc_current_gain;
        float time_constant;
        
        if (gain_diff > 0) {
            // 增益增加，使用释放时间
            time_constant = processing->agc_release_time;
        } else {
            // 增益减少，使用攻击时间
            time_constant = processing->agc_attack_time;
        }
        
        // 一阶低通滤波器平滑增益
        float alpha = 1.0f - expf(-1.0f / (time_constant * LITTLE_TEST_SAMPLE_RATE / LITTLE_TEST_FRAME_SIZE));
        processing->agc_current_gain += alpha * gain_diff;
        
        // 应用增益
        uint32_t i;
        for (i = 0; i < count; i++) {
            float sample_f = (float)samples[i] * processing->agc_current_gain;
            
            // 限幅
            if (sample_f > 32767.0f) {
                sample_f = 32767.0f;
            } else if (sample_f < -32768.0f) {
                sample_f = -32768.0f;
            }
            
            samples[i] = (int16_t)sample_f;
        }
    }
}

/**
 * @brief 噪声门限处理
 * @param processing 音频预处理结构体指针
 * @param samples 音频样本数据
 * @param count 样本数量
 */
static void apply_noise_gate(audio_processing_t *processing, int16_t *samples, uint32_t count)
{
    if (!processing->noise_gate_enabled) {
        return;
    }
    
    // 计算当前帧的RMS电平
    uint32_t current_rms = calculate_rms_level(samples, count);
    
    // 如果电平低于门限，将信号静音
    if (current_rms < processing->noise_gate_threshold) {
        memset_s(samples, count * sizeof(int16_t), 0, count * sizeof(int16_t));
    }
}

/**
 * @brief 处理一帧音频数据
 * @param processing 音频预处理结构体指针
 * @param input_frame 输入的音频帧（双声道，但只有一个声道有数据）
 * @param output_frame 输出的音频帧（单声道）
 * @return 0成功，其他失败
 */
int audio_processing_process_frame(audio_processing_t *processing, 
                                  const audio_frame_t *input_frame, 
                                  audio_frame_t *output_frame)
{
    if (processing == NULL || !processing->initialized || 
        input_frame == NULL || output_frame == NULL) {
        osal_printk("[AUDIO_PROC] Invalid parameters\n");
        return -1;
    }

    if (input_frame->size != LITTLE_TEST_FRAME_SIZE * 2 * sizeof(int16_t)) {
        osal_printk("[AUDIO_PROC] Invalid input frame size: %d\n", input_frame->size);
        return -1;
    }

    // 复制输入数据到原始缓冲区
    memcpy_s(processing->raw_buffer, 
             LITTLE_TEST_FRAME_SIZE * 2 * sizeof(int16_t),
             input_frame->data, 
             input_frame->size);

    // 提取单声道数据（假设左声道有数据，右声道为0）
    uint32_t i;
    for (i = 0; i < LITTLE_TEST_FRAME_SIZE; i++) {
        // 从双声道数据中提取左声道
        processing->temp_buffer[i] = processing->raw_buffer[i * 2];
    }

    // 应用高通滤波器去除直流分量
    for (i = 0; i < LITTLE_TEST_FRAME_SIZE; i++) {
        processing->temp_buffer[i] = high_pass_filter(processing, processing->temp_buffer[i]);
    }

    // 复制到处理缓冲区
    memcpy_s(processing->processed_buffer, 
             LITTLE_TEST_FRAME_SIZE * sizeof(int16_t),
             processing->temp_buffer, 
             LITTLE_TEST_FRAME_SIZE * sizeof(int16_t));

    // 应用噪声门限
    apply_noise_gate(processing, processing->processed_buffer, LITTLE_TEST_FRAME_SIZE);

    // 应用自动增益控制
    apply_agc(processing, processing->processed_buffer, LITTLE_TEST_FRAME_SIZE);

    // 更新统计信息
    processing->peak_level = find_peak_level(processing->processed_buffer, LITTLE_TEST_FRAME_SIZE);
    processing->rms_level = calculate_rms_level(processing->processed_buffer, LITTLE_TEST_FRAME_SIZE);
    processing->frame_count++;

    // 设置输出帧
    output_frame->data = processing->processed_buffer;
    output_frame->size = LITTLE_TEST_FRAME_SIZE * sizeof(int16_t);
    output_frame->timestamp = input_frame->timestamp;
    output_frame->sequence = input_frame->sequence;

    // 每100帧打印一次统计信息
    if (processing->frame_count % 100 == 0) {
        osal_printk("[AUDIO_PROC] Frame %d: Peak=%d, RMS=%d, AGC_Gain=%.2f\n", 
                    processing->frame_count, processing->peak_level, 
                    processing->rms_level, processing->agc_current_gain);
    }

    return 0;
}

/**
 * @brief 设置噪声门限
 * @param processing 音频预处理结构体指针
 * @param threshold 噪声门限值
 * @param enabled 是否启用噪声门限
 * @return 0成功，其他失败
 */
int audio_processing_set_noise_gate(audio_processing_t *processing, uint32_t threshold, bool enabled)
{
    if (processing == NULL || !processing->initialized) {
        return -1;
    }

    processing->noise_gate_threshold = threshold;
    processing->noise_gate_enabled = enabled;
    
    osal_printk("[AUDIO_PROC] Noise gate: threshold=%d, enabled=%s\n", 
                threshold, enabled ? "true" : "false");
    
    return 0;
}

/**
 * @brief 设置AGC参数
 * @param processing 音频预处理结构体指针
 * @param target_level 目标电平
 * @param max_gain 最大增益
 * @param min_gain 最小增益
 * @return 0成功，其他失败
 */
int audio_processing_set_agc(audio_processing_t *processing, 
                            uint32_t target_level, 
                            float max_gain, 
                            float min_gain)
{
    if (processing == NULL || !processing->initialized) {
        return -1;
    }

    if (max_gain < min_gain || min_gain <= 0 || max_gain <= 0) {
        osal_printk("[AUDIO_PROC] Invalid AGC parameters\n");
        return -1;
    }

    processing->agc_target_level = target_level;
    processing->agc_max_gain = max_gain;
    processing->agc_min_gain = min_gain;
    
    osal_printk("[AUDIO_PROC] AGC: target=%d, max_gain=%.2f, min_gain=%.2f\n", 
                target_level, max_gain, min_gain);
    
    return 0;
}

/**
 * @brief 获取处理统计信息
 * @param processing 音频预处理结构体指针
 * @param peak_level 输出峰值电平
 * @param rms_level 输出RMS电平
 * @param current_gain 输出当前增益
 * @return 0成功，其他失败
 */
int audio_processing_get_stats(audio_processing_t *processing, 
                              uint32_t *peak_level, 
                              uint32_t *rms_level, 
                              float *current_gain)
{
    if (processing == NULL || !processing->initialized) {
        return -1;
    }

    if (peak_level != NULL) {
        *peak_level = processing->peak_level;
    }
    
    if (rms_level != NULL) {
        *rms_level = processing->rms_level;
    }
    
    if (current_gain != NULL) {
        *current_gain = processing->agc_current_gain;
    }
    
    return 0;
}